import os
import sys

import pandas as pd

# 获取模型
current_path = sys.path[0]
DATA_FILE_PATH = os.path.join(current_path, 'file\\')


class DataUtil:
    def __init__(self, train_data):
        self.train_data = pd.read_csv(DATA_FILE_PATH + train_data, encoding="utf-8")
        self.basic_data = {
            "houseNumberSuffix": {"号", "－", "号－", "号之", "号西", "号楼", "区", "幢", "队", "号西", "号上", "号下", "组"},  # 门牌后缀
            "buildingNumberSuffix": {"幢", "栋", "排", "舍", "楼", "号楼", "弄", "门", "队"},  # 栋楼后缀
            "roomNumberSuffix": {"层", "室", "号"},
            "JLX": pd.read_csv(DATA_FILE_PATH + "基础字段字典项.csv", encoding="utf-8").values[0:, 1]
        }

    def get_column(self, column_name):
        return self.train_data[column_name]

    def get_column_names(self):
        return self.train_data.keys()
